import psycopg2
from langchain_core.stores import BaseStore

# PostgreSQL 连接配置
PG_HOST = "192.168.6.20"
PG_PORT = 5432
PG_USER = "ipp_air"
PG_PASSWORD = "Air@yutu123."
PG_DB = "yutu_ipp_air"


def connect_pg():
    """连接 PostgreSQL"""
    conn = psycopg2.connect(
        host=PG_HOST, port=PG_PORT, user=PG_USER, password=PG_PASSWORD, dbname=PG_DB
    )
    return conn


class PGStore(BaseStore):
    """适配 PostgreSQL 作为 Parent Document 存储"""

    def __init__(self):
        self.conn = psycopg2.connect(
            host=PG_HOST, port=PG_PORT, user=PG_USER, password=PG_PASSWORD, dbname=PG_DB
        )

    def mget(self, keys):
        """获取多个 Parent 文档（适用于 Knowledge）"""
        if not keys:
            return {}
        cur = self.conn.cursor()
        format_strings = ",".join(["%s"] * len(keys))
        query = f"SELECT parent_id, content FROM ipp_parent_docs WHERE parent_id IN ({format_strings})"
        cur.execute(query, tuple(keys))
        results = {row[0]: row[1] for row in cur.fetchall()}
        cur.close()
        return results

    def get(self, key):
        """获取单个 Parent 文档"""
        cur = self.conn.cursor()
        cur.execute("SELECT content FROM ipp_parent_docs WHERE parent_id = %s", (key,))
        result = cur.fetchone()
        cur.close()
        return result[0] if result else None

    def set(self, key, value):
        """插入/更新 Parent 文档"""
        cur = self.conn.cursor()
        cur.execute(
            """
            INSERT INTO ipp_parent_docs (parent_id, content) 
            VALUES (%s, %s) 
            ON CONFLICT (parent_id) DO UPDATE SET content = EXCLUDED.content
        """,
            (key, value),
        )
        self.conn.commit()
        cur.close()

    def mset(self, keys, values):
        """批量插入/更新 Parent 文档"""
        cur = self.conn.cursor()
        for key, value in zip(keys, values):
            cur.execute(
                """
                INSERT INTO ipp_parent_docs (parent_id, content) 
                VALUES (%s, %s) 
                ON CONFLICT (parent_id) DO UPDATE SET content = EXCLUDED.content
            """,
                (key, value),
            )
        self.conn.commit()
        cur.close()

    def mdelete(self, keys):
        """批量删除 Parent 文档"""
        cur = self.conn.cursor()
        format_strings = ",".join(["%s"] * len(keys))
        query = f"DELETE FROM ipp_parent_docs WHERE parent_id IN ({format_strings})"
        cur.execute(query, tuple(keys))
        self.conn.commit()
        cur.close()

    def yield_keys(self):
        """遍历所有 Parent 文档的 ID"""
        cur = self.conn.cursor()
        cur.execute("SELECT parent_id FROM ipp_parent_docs")
        for row in cur.fetchall():
            yield row[0]
        cur.close()


def insert_into_pg(parent_docs):
    """存储 Parent 文档到 PostgreSQL"""
    conn = connect_pg()
    cur = conn.cursor()

    for doc in parent_docs:
        cur.execute(
            """
            INSERT INTO ipp_parent_docs (parent_id, title, content) 
            VALUES (%s, %s, %s)
        """,
            (doc["id"], doc["title"], doc["content"]),
        )

    conn.commit()  # ✅ 确保事务提交
    cur.close()
    conn.close()
    print(f"✅ 存入 {len(parent_docs)} 个父文档到 PostgreSQL")
